import sys

from PyQt5.QtCore import Qt, QCoreApplication
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog
import lib.script
from UI import Ui_MainWindow as JudgeWindow
from MainUI import Ui_MainWindow as MainWindow
from PaintUI import Ui_MainWindow as PaintWindow
from Paint import Winform
from lib.pix2pix.test import birth


class Main(QMainWindow, MainWindow):
    def __init__(self, parent=None):
        super(Main, self).__init__()
        self.setupUi(self)
        self.JudgeButton.clicked.connect(self.openJudge)
        self.TeachButton.clicked.connect(self.openPaint)
        self.exitButton.clicked.connect(QCoreApplication.quit)

    def openJudge(self):
        self.JudgeWindow = Judge()
        self.JudgeWindow.show()
        self.close()

    def openPaint(self):
        self.PaintWindow = Paint()
        self.PaintWindow.show()
        self.close()


class Paint(QMainWindow, PaintWindow):
    def __init__(self, parent=None):
        super(Paint, self).__init__()
        self.setupUi(self)
        self.exitButton.clicked.connect(self.reback)
        self.confirmButton.clicked.connect(self.openPrint)
        self.showButton.clicked.connect(self.showTemp)
        self.birthButton.clicked.connect(self.birth_show)

    def openPrint(self):
        self.PaintWindow = Winform()
        self.PaintWindow.show()

    def showTemp(self):
        self.image_view.clear()
        self.image = QImage('lib/temp/test/temp.jpg').scaled(256, 256)
        self.image_view.setPixmap(QPixmap.fromImage(self.image))

    def reback(self):
        self.close()
        self.MainWindow = Main()
        self.MainWindow.show()

    def birth_show(self):
        birth()
        self.image_view.clear()
        self.image = QImage('lib/results/images/temp_fake_B.jpg').scaled(561, 471)
        self.image_view.setPixmap(QPixmap.fromImage(self.image))


class Judge(QMainWindow, JudgeWindow):
    def __init__(self, parent=None):
        super(Judge, self).__init__()
        self.setupUi(self)
        self.signal()

    def imageGet(self, checked):
        self.filename = QFileDialog.getOpenFileName(self, "OpenFile", ".", "Image Files(*.jpg *.jpeg *.png)")[0]
        if len(self.filename):
            self.image = QImage(self.filename).scaled(420, 241)
            self.image_view.setPixmap(QPixmap.fromImage(self.image))

    def classfyGet(self, checked):
        text = lib.script.GetClassfy(self.filename)
        self.text_view.setText('Result:  ' + text)
        if text is None:
            return
        self.text_view.show()
        lib.script.Delete()

    def signal(self):
        self.input_button.clicked.connect(self.imageGet)
        self.ok_button.clicked.connect(self.classfyGet)
        self.report_button.clicked.connect(self.reback)

    def reback(self):
        self.close()
        self.MainWindow = Main()
        self.MainWindow.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = Main()
    window.setWindowFlags(Qt.WindowCloseButtonHint)
    window.show()
    sys.exit(app.exec_())
